#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""计算斜率
给定 DEM 的栅格文件, 输出对应斜率栅格。

Example:
    $ python slope.py input.tif out.tif
"""

import os
import sys
import cv2
import numpy as np
from osgeo import gdal

def compute_slope(srcband, x_cellsize, y_cellsize):
    print(f"x_cellsize: {x_cellsize}, y_cellsize: {y_cellsize}")
    # print(f"RasterXSize: {gtif.RasterXSize}, RasterYSize: {gtif.RasterYSize}")
    raster_x_size = srcband.XSize
    raster_y_size = srcband.YSize
    # print(f"XSize: {raster_x_size}, YSize: {raster_y_size}")

    raster_array = srcband.ReadAsArray(0, 0, raster_x_size, raster_y_size)
    # print(srcband.GetNoDataValue())

    sobel_x = cv2.Sobel(raster_array, cv2.CV_64F, 1, 0)
    sobel_y = cv2.Sobel(raster_array, cv2.CV_64F, 0, 1)

    slope = (np.arctan(np.sqrt((sobel_x / 8 / x_cellsize) ** 2 + (sobel_y / 8 / y_cellsize) ** 2))) * 180 / np.pi

    return slope

def main(input_file, out_tif='out.tif'):
    gtif = gdal.Open(input_file)

    if gtif is None:
        print(f"Unable to open {input_file}")
        sys.exit(1)

    try:
        srcband = gtif.GetRasterBand(1)
    except RuntimeError as e:
        print(f"No band found")
        print(e)
        sys.exit(1)

    geo_transform = gtif.GetGeoTransform()
    no_data_value = srcband.GetNoDataValue()
    slope = compute_slope(srcband, geo_transform[1], geo_transform[5])
    # print(slope.dtype)

    driver = gdal.GetDriverByName('GTiff')
    driver.Register()

    out_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), out_tif)
    out_tif = driver.Create(out_file, srcband.XSize, srcband.YSize, 1, gdal.GDT_Float32)
    out_tif.SetGeoTransform(gtif.GetGeoTransform())
    out_tif.SetProjection(gtif.GetProjection())

    out_band = out_tif.GetRasterBand(1)
    out_band.WriteArray(slope, 0, 0)
    out_band.SetNoDataValue(no_data_value)
    out_band.FlushCache()

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print("""[ ERROR ] you must supply at least one arguments:
1) input raster 2) output raster (options)""")
        sys.exit(1)
    elif len(sys.argv) == 2:
        main(sys.argv[1])
    else:
        main(sys.argv[1], sys.argv[2])
